from flask import Flask, jsonify
from datetime import datetime
import pandas as pd
from gensim.models import Word2Vec
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import accuracy_score, recall_score

app = Flask(__name__)

# Function to train the model
def train_model():
    # 数据读取和预处理
    train_df = pd.read_csv('train_product.csv')
    test_df = pd.read_csv('test_product.csv')

    # 提取文本描述列
    train_descriptions = train_df['productDescription']
    test_descriptions = test_df['productDescription']

    # 对描述进行分词
    train_descriptions = train_descriptions.apply(lambda x: x.split())
    test_descriptions = test_descriptions.apply(lambda x: x.split())

    # 合并所有描述以便训练Word2Vec模型
    all_descriptions = pd.concat([train_descriptions, test_descriptions])

    # 训练Word2Vec模型
    word2vec_model = Word2Vec(all_descriptions, vector_size=100, window=5, min_count=1, workers=4)

    # 将描述转换为特征向量
    def get_feature_vector(description, model):
        words = description
        feature_vector = np.mean([model.wv[word] for word in words if word in model.wv], axis=0)
        return feature_vector

    train_vectors = train_descriptions.apply(lambda x: get_feature_vector(x, word2vec_model))
    test_vectors = test_descriptions.apply(lambda x: get_feature_vector(x, word2vec_model))

    # 构建和训练深度神经网络模型
    class CategoryModel(nn.Module):
        def __init__(self, input_dim, hidden_dim, output_dim):
            super(CategoryModel, self).__init__()
            self.fc1 = nn.Linear(input_dim, hidden_dim)
            self.relu = nn.ReLU()
            self.fc2 = nn.Linear(hidden_dim, output_dim)

        def forward(self, x):
            out = self.fc1(x)
            out = self.relu(out)
            out = self.fc2(out)
            return out

    # 参数设置
    input_dim = 100  # Word2Vec向量的维度
    hidden_dim = 64
    output_dim = train_df['category'].nunique()  # 类别数量
    num_epochs = 20
    learning_rate = 0.001

    # 数据准备
    train_X = np.stack(train_vectors.values)
    train_y = train_df['category'].values
    test_X = np.stack(test_vectors.values)
    test_y = test_df['category'].values

    # 转换为PyTorch张量
    train_X = torch.tensor(train_X, dtype=torch.float32)
    train_y = torch.tensor(train_y, dtype=torch.long)
    test_X = torch.tensor(test_X, dtype=torch.float32)
    test_y = torch.tensor(test_y, dtype=torch.long)

    # 模型实例化
    model = CategoryModel(input_dim, hidden_dim, output_dim)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 模型训练
    for epoch in range(num_epochs):
        model.train()
        outputs = model(train_X)
        loss = criterion(outputs, train_y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

    # 保存模型
    torch.save(model.state_dict(), 'categoryModel.pth')

    # 模型测试
    model.eval()
    with torch.no_grad():
        test_outputs = model(test_X)
        _, predicted = torch.max(test_outputs, 1)

    # 计算指标
    accuracy = accuracy_score(test_y, predicted)
    recall = recall_score(test_y, predicted, average='weighted')

    # 打印结果
    print(f'Accuracy: {accuracy:.4f}')
    print(f'Recall: {recall:.4f}')

    # 保存Word2Vec模型
    word2vec_model.save('word2vec_model.model')

    return {
        'accuracy': accuracy,
        'recall': recall
    }

@app.route('/productCategoryTrain', methods=['POST'])
def train():
    results = train_model()
    return jsonify(results)

if __name__ == '__main__':
    app.run(debug=True)
